Problem

A non-finite box coordinate makes every comparison in the "inverse elements are out of feature map boundary" check false, so the sample point is neither rejected by the empty branch nor clamped by the y <= 0 / x <= 0 clamps. The index is then computed from (int)y with y == NaN, which is undefined behaviour and yields INT_MIN on x86. pos1..pos4 become large negative offsets and the forward kernels read far outside the input buffer:

import torch, torchvision
nan = float("nan")
x = torch.zeros(1, 1, 4, 4)
b = torch.tensor([[0.0, nan, nan, nan, nan]])
torchvision.ops.roi_align(x, b, (1, 1), 1.0, 1, True)   # Segmentation fault
torch.ops.torchvision.ps_roi_align(x, b, 1.0, 1, 1, 1)  # Segmentation fault

UBSan and ASan on the unmodified kernel:

roi_align_common.h:83:28: runtime error: nan is outside the range of representable values of type 'int'
roi_align_common.h:109:27: runtime error: signed integer overflow: -2147483648 * 4 cannot be represented in type 'int'
AddressSanitizer: SEGV on unknown address 0x505e00000020 ... caused by a READ memory access

It is not only NaN that gets there, and an input-side isfinite filter would not be enough. Measured on x = torch.zeros(1, 1, 4, 4), roi_align(..., (1, 1), 1.0, 1, True):

The backward kernels reach the same undefined conversion but do not corrupt memory: every write is already gated on x_low >= 0 && x_high >= 0 && y_low >= 0 && y_high >= 0, which INT_MIN fails. So this is an out-of-bounds read only, and because all four interpolation weights are NaN whenever the coordinate is, whatever is read is multiplied by NaN and cannot be observed in the output. The impact is a crash, not disclosure.

Fix

Write the guard as !(in range) instead of (out of range) so a NaN coordinate takes the empty branch. For every finite input the two forms are exactly complementary, so this is not a behaviour change outside the non-finite case.

The same guard is duplicated at 10 sites in 6 files, and all are updated:

torchvision/csrc/ops/cpu/roi_align_common.h:60
torchvision/csrc/ops/cpu/roi_align_kernel.cpp:133
torchvision/csrc/ops/cpu/ps_roi_align_kernel.cpp:26,169
torchvision/csrc/ops/cuda/roi_align_kernel.cu:35,173
torchvision/csrc/ops/cuda/ps_roi_align_kernel.cu:35,170
torchvision/csrc/ops/mps/mps_stable_kernels.h:60,167

Verification

Built torchvision from source (v0.23.0, same kernels, against torch 2.8.0) with and without this change and ran each case in its own process:

Six of six crash before, six of six clean after. Note the regression test kills the interpreter on unpatched code rather than reporting an assertion failure, as memory-safety tests do.

Other checks:

• No behaviour change for finite inputs. Compiling the old and new roi_align_common.h side by side in one binary and comparing the resulting PreCalc structs bitwise over 20,055,768 sample points (random geometries, spatial scales 1e-3/1/1e6, 12,865,481 of them taking the empty branch) gives 0 mismatches.
• Existing tests. pytest test/test_ops.py -k "TestRoIAlign or TestPSRoIAlign or TestRoIPool or TestPSRoIPool" on the patched build: 143 passed, 0 failed.
• Sanitizers. All three coordinate classes go from float-cast-overflow + ASan SEGV to a clean 0.0 result.
• CPU and GPU now agree. CUDA/HIP never crashed here, because AMD and NVIDIA float→int conversion returns 0 for NaN rather than INT_MIN (measured: (int)NaNf is 0 on gfx1201, INT_MIN on x86-64). The GPU therefore sampled index 0 and returned nan where the CPU faulted. After this change both return 0, which is what an out-of-range sample point has always produced. This is the one user-visible output change, and only for input that was undefined behaviour before.
• New test_non_finite_boxes in RoIOpTester, so it runs for roi_align, ps_roi_align, roi_pool and ps_roi_pool, forward and backward, on every device: 12 passed on CPU (the pool ops were already safe; they are covered so the whole family is pinned).

Not fixed here

int roi_batch_ind = offset_rois[0]; in the same kernels is an unchecked conversion too, and is not range-checked against the batch size anywhere (24 sites). The same NaN payload placed in column 0 still reads out of bounds at some shapes, and so does a plain out-of-range integer index — that is #4828, open since 2021. It needs a host-side check rather than a kernel-side predicate, so it is left for a separate change.

Disclosure per AI_POLICY.md: the analysis and this patch were prepared with AI assistance (Claude). Every measurement quoted above was produced by actually building and running the code described, not inferred.